from pretreatment import Img2Array
import numpy as np
from fullyConnect import FullyConnect
from sigmoid import Sigmoid
from data import Data
from img_pre_deal import ImagePre
from code_add import Code
import sys,os
sys.path.append(os.path.dirname(__file__))
class Predict:
    def __init__(self,imgPath=os.path.join(os.getcwd(),'sourceImg','code.png')) -> None:
        self.inner_layers=[]
        self.inner_layers_symbol=[]
        self.f1=FullyConnect(15*25,4)
        self.f2=FullyConnect(4,9)
        self.f3=FullyConnect(10*25,2)
        self.base=os.path.dirname(__file__)
        self.imgPath=imgPath
        self.initLayer()

    def initLayer(self):
        self.inner_layers.append(self.f1)
        self.inner_layers.append(Sigmoid())
        self.inner_layers.append(self.f2)
        self.inner_layers.append(Sigmoid())      
        args1_w=np.array(np.loadtxt(os.path.join(self.base,'args','args1_w.txt')))
        args1_b=np.array(np.loadtxt(os.path.join(self.base,'args','args1_b.txt')))
        args2_w=np.array(np.loadtxt(os.path.join(self.base,'args','args2_w.txt')))
        args2_b=np.array(np.loadtxt(os.path.join(self.base,'args','args2_b.txt')))
        args1_b=args1_b.reshape((args1_b.size,1))
        args2_b=args2_b.reshape((args2_b.size,1))
        self.f1.weights,self.f1.bias=(args1_w,args1_b)
        self.f2.weights,self.f2.bias=(args2_w,args2_b)

        self.inner_layers_symbol.append(self.f3)
        self.inner_layers_symbol.append(Sigmoid())
        args3_w=np.array(np.loadtxt(os.path.join(self.base,'args','args3_w.txt')))
        args3_b=np.array(np.loadtxt(os.path.join(self.base,'args','args3_b.txt')))
        args3_b=args3_b.reshape((args3_b.size,1))
        self.f3.weights,self.f3.bias=(args3_w,args3_b)

    def preDeal(self):
        originDir=os.path.dirname(self.imgPath)
        targetDir=os.path.join(self.base,'tmp','noiseImg') 
        self.del_file(os.path.join(self.base,'tmp'))
        imgpre=ImagePre(originDir,targetDir)
        imgpre.main()
    
    def del_file(self,path):
        ls = os.listdir(path)
        for i in ls:
            c_path = os.path.join(path, i)
            if os.path.isdir(c_path):
                self.del_file(c_path)
            else:
                os.remove(c_path)
    
    def preDict(self):
        img=Img2Array()
        img.main(os.path.join(self.base,'data','predict.txt'),os.path.join(self.base,'data','predict.npy'))
        datalayer=Data(os.path.join(self.base,'data','predict.npy'),2)        
        data,_=datalayer.forward()
        x=data[0]
        for layer in self.inner_layers:
            x=layer.forward(x)
        r=np.argmax(x,axis=1)+1
        print("debug",r)
        return r[0][0],r[1][0]
        # print('pridict:',np.argmax(x,axis=1)+1) #预测值
        # symbolIndex=self.preDictSymbol()
        # print(symbolIndex)
        # symbols=['+','*']
        # # return np.argmax(x,axis=1)+1,symbols[int(symbolIndex)]

    def preDictSymbol(self):
        img=Img2Array()
        img.main(os.path.join(self.base,'data','predictSymbol.txt'),os.path.join(self.base,'data','predictSymbol.npy'))
        datalayer=Data(os.path.join(self.base,'data','predictSymbol.npy'),2)        
        data,_=datalayer.forward()
        x=data[0]
        for layer in self.inner_layers_symbol:
            x=layer.forward(x)
        symbols=['+','*']
        index=np.argmax(x,axis=1)[0][0]
        return (symbols[int(index)],)
    
    def main(self):
        self.preDeal()
        r1=self.preDict()
        return r1+self.preDictSymbol()

    def getCode(self,base64Str):
        code=Code()
        code.base64_to_image(base64Str)
        self.preDeal()
        r1=self.preDict()
        r=r1+self.preDictSymbol()
        print(r)
        symbol=r[2]
        c1=int(r[0])
        c2=int(r[1])
        result=0
        if symbol=='+':
            result=c1+c2
        elif symbol=='*':
            result=c1*c2
        return result
    def getCode(self):
        r=self.main()
        print(r)
        symbol=r[2]
        c1=int(r[0])
        c2=int(r[1])
        result=0
        if symbol=='+':
            result=c1+c2
        elif symbol=='*':
            result=c1*c2
        return result

if __name__ == '__main__':    
    # 直接根据base64预测
    # base64str=""
    # p=Predict() 
    # print(p.getCode(base64Str=base64str))
    
    # 根据图片地址预测，每次预测时会将sourceImg下的所有图片预处理,所以每次测试时，需要保证sourceImg下只有一张图片才会是正确的答案
    imgpath=os.path.join(os.path.dirname(__file__),'sourceImg','code1.png') 
    p2=Predict(imgPath=imgpath)
    print(p2.getCode())
